# -*- coding: utf-8 -*-
"""
Created on Fri Sep 11 13:23:00 2020

@author: Tom
"""
import os.path
import os
import time
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import random
import tensorflow as tf
from midas import Midas

from sklearn.preprocessing import MinMaxScaler

np.random.seed(89)
random.seed(89)

path = 'kropko/anes_midas/'

categorical = ['religion','married', 'vote']

binary = ['imp_enviro','female', 'white']
  
for k in range(1,21):
  data_0 = pd.read_csv(path+"kropko_ANES_data_missing_"+str(k)+".csv")

  data_1 = data_0[categorical]
  data_0.drop(categorical, axis= 1, inplace= True)

  constructor_list = [data_0]
  columns_list = []

  for column in data_1.columns:
    na_temp = data_1[column].isnull()
    temp = pd.get_dummies(data_1[column], prefix = column)
    temp[na_temp] = np.nan
    constructor_list.append(temp)
    columns_list.append(list(temp.columns.values))

  data_0 = pd.concat(constructor_list, axis=1)

  na_loc = data_0.isnull()
  data_0[na_loc] = np.nan

  scaler = MinMaxScaler()

  data_0_scaled = scaler.fit_transform(data_0)
  data_0_scaled = pd.DataFrame(data_0_scaled, columns = data_0.columns)


  imp = Midas(layer_structure = [512,512,512], vae_layer = False, seed = 89)
  
  imp.build_model(data_0_scaled,
                  categorical_columns = binary, softmax_columns = columns_list)

  imp.train_model(training_epochs=20, verbose=True)

  imputations = imp.generate_samples(m=5).output_list
  
  n = 1
  for i in imputations:
    file_name = path + "ANES_imp_mid_" + str(k) + "_" + str(n) + ".csv"
    out_file = i[data_0_scaled.columns]
    out_file = scaler.inverse_transform(out_file)
    out_file = pd.DataFrame(out_file, columns = data_0_scaled.columns)
    out_file.to_csv(file_name, index=False)
    n +=1
